from common_import import *


def calculate_total_profit(profit_csv, planting_csv):
    # 读取CSV文件
    profit_df = pd.read_csv(profit_csv)
    planting_df = pd.read_csv(planting_csv)

    # 为每个地块添加地块类型
    planting_df["地块类型"] = planting_df["种植地块"].apply(mapping.get_land_type)
    normal_greenhouse_first_season = profit_df[
        (profit_df["FieldType"] == "普通大棚 ")
        & (profit_df["PlantingSeason"] == "第一季")
    ].copy()

    # 修改这部分数据的FieldType为智慧大棚
    normal_greenhouse_first_season["FieldType"] = "智慧大棚"

    # 将智慧大棚的第一季数据添加到利润表
    profit_df = pd.concat(
        [profit_df, normal_greenhouse_first_season], ignore_index=True
    )
    print(profit_df)
    # 根据作物编号、地块类型和种植季次来合并数据
    merged_df = pd.merge(
        planting_df,
        profit_df,
        how="left",
        left_on=["作物编号", "地块类型", "种植季次"],
        right_on=["CropID", "FieldType", "PlantingSeason"],
    )

    # 计算每个地块的总利润 = 种植面积 * 每亩利润
    merged_df["总利润"] = merged_df["种植面积/亩"] * merged_df["Profit_per_Mu"]
    # 按地块类型和种植季次汇总总利润
    result = merged_df.groupby(["地块类型", "种植季次"])["总利润"].sum().reset_index()
    return result


if __name__ == "__main__":
    result = calculate_total_profit("data/利润.csv", "data/2_每个地种什么.csv")
    print(result)
    print(result["总利润"].sum())
